package org.teresoft.collections;

import static java.lang.Integer.bitCount;
import static java.lang.System.arraycopy;

import org.teresoft.Preconditions;

@SuppressWarnings({"unchecked", "rawtypes"})
public class HashTrie<K, V> implements Map<K, V> {

    private static final HashTrie EMPTY = new HashTrie(null, 0);
    
    private final Node node;
    private final int size;
    
    private HashTrie(Node node, int size) {
        this.node = node;
        this.size = size;
    }
    
    public static <K, V> HashTrie<K, V> create() {
        return (HashTrie<K, V>) EMPTY;
    }

    @Override
    public V get(Object key) {
        
        Preconditions.checkNotNull(key, "'key' must not be null");
        
        if (node == null) {
            return null;
        }
        
        return (V) node.get(hash(key.hashCode()), key);
    }

    @Override
    public HashTrie<K, V> put(K key, V value, Out<V> holder) {
        
        Preconditions.checkNotNull(key, "'key' must not be null");
        Preconditions.checkNotNull(value, "'value' must not be null");
        Preconditions.checkNotNull(holder, "'holder' must not be null");
        
        Node newNode;
        int hash = hash(key.hashCode());
        if (node != null) {
            newNode = node.put(0, hash, key, value, (Out<Object>) holder);
        } else {
            newNode = new LeafNode(hash, key, value);
        }
        
        int newSize = (holder.value == null) ? size + 1 : size;
        
        return new HashTrie<K, V>(newNode, newSize);
    }

    @Override
    public HashTrie<K, V> remove(K key, Out<V> holder) {
        
        Preconditions.checkNotNull(key, "'key' must not be null");
        Preconditions.checkNotNull(holder, "'holder' must not be null");
        
        if (node == null) {
            return this;
        }
        
        int hash = hash(key.hashCode());
        Node newNode = node.remove(hash, key, (Out<Object>) holder);        
        
        if (node == newNode) {
            return this;
        }
        
        return new HashTrie<K, V>(newNode, size - 1);
    }

    /**
     * Stolen from java.util.HashMap.
     * 
     * @param h
     * @return
     */
    static int hash(int h) {
        // This function ensures that hashCodes that differ only by
        // constant multiples at each bit position have a bounded
        // number of collisions (approximately 8 at default load factor).
        h ^= (h >>> 20) ^ (h >>> 12);
        return h ^ (h >>> 7) ^ (h >>> 4);
    }
    
    interface Node {
        Object get(final int hashCode, Object key);
        Node put(final int level, final int hashCode, Object key, Object value, Out<Object> out);
        Node remove(final int hashCode, Object key, Out<Object> out);
    }
    
    static final class BitMappedNode implements Node {
        
        private final int level;
        private final Node[] table;
        private final int bitMap;
        
        private BitMappedNode(int level, Node[] table, int bitMap) {
            this.level = level;
            this.table = table;
            this.bitMap = bitMap;
        }
        
        @Override
        public Object get(final int hashCode, Object key) {
            int bitIndex = 1 << getPartialHash(hashCode, level);
            if ((bitIndex & bitMap) != 0) {
                int index = bitCount((bitIndex - 1) & bitMap);
                return table[index].get(hashCode, key);
            }
            return null;
        }
        
        @Override
        public BitMappedNode put(int level, final int hashCode, Object key, Object value, Out<Object> out) {
            int bitIndex = 1 << getPartialHash(hashCode, level);
            int index = bitCount((bitIndex -1) & bitMap);
            
            if ((bitIndex & bitMap) == 0) {
                LeafNode node = new LeafNode(hashCode, key, value);
                return new BitMappedNode(level, copyAndInsert(table, node, index), bitMap | bitIndex); 
            } else {
                Node node = table[index].put(level + 1, hashCode, key, value, out);
                return new BitMappedNode(level, copyAndReplace(table, node, index), bitMap); 
            }
        }
        
        @Override
        public Node remove(final int hashCode, Object key, Out<Object> out) {
            int bitIndex = 1 << getPartialHash(hashCode, level);
            
            if ((bitIndex & bitMap) != 0) {
                int index = bitCount((bitIndex - 1) & bitMap);
                Node child = table[index].remove(hashCode, key, out);
                if (out.value != null) {
                    if (child == null) {
                        if (table.length != 2) {                    
                            return new BitMappedNode(level, copyAndRemove(table, index), bitMap ^ bitIndex);
                        } else {
                            return table[index ^ 1];
                        }
                    } else {
                        return new BitMappedNode(level, copyAndReplace(table, child, index), bitMap);
                    }
                }
            }
            
            return this;
        }
        
        private static int getPartialHash(final int hashCode, int level) {
            int levelShift = level * 5;
            int mask = 0x1F << levelShift;
            return ((mask & hashCode) >>> levelShift);
        }

        private static BitMappedNode create(int level, int hashCodeA, Node nodeA, int hashCodeB, Node nodeB) {
            int partialHashA = getPartialHash(hashCodeA, level);
            int partialHashB = getPartialHash(hashCodeB, level);
            
            Node[] table;
            if (partialHashA < partialHashB) {
                table = new Node[] { nodeA, nodeB };
            } else if (partialHashA > partialHashB) {
                table = new Node[] { nodeB, nodeA };
            } else if (hashCodeA != hashCodeB) {
                Node node = create(level + 1, hashCodeA, nodeA, hashCodeB, nodeB);
                table = new Node[] { node };
            } else {
                throw new UnsupportedOperationException();
            }
            
            return new BitMappedNode(level, table, (1 << partialHashA) | (1 << partialHashB));
        }
    }
        
    static final class LeafNode implements Node {
        
        private final int hashCode;
        private final Object key;
        private final Object value;

        public LeafNode(final int hashCode, Object key, Object value) {
            this.hashCode = hashCode;
            this.key = key;
            this.value = value;
        }

        @Override
        public Object get(final int hashCode, Object key) {
            if (keyMatches(hashCode, key)) {
                return value;
            }
            return null;
        }

        @Override
        public Node put(int level, final int hashCode, Object key, Object value, Out<Object> out) {
            Node node = new LeafNode(hashCode, key, value);
            if (this.hashCode == hashCode) {
                if (this.key == key || this.key.equals(key)) {
                    out.value = this.value;
                    return node;
                } else {
                    return ListNode.create(this, node);
                }
            } else {
                return BitMappedNode.create(level, this.hashCode, this, hashCode, node);
            }
        }
        
        @Override
        public Node remove(final int hashCode, Object key, Out<Object> out) {
            if (keyMatches(hashCode, key)) {
                out.value = this.value;
                return null;
            } else {
                return this;
            }
        }
        
        private boolean keyMatches(final int hashCode, Object key) {
            return this.hashCode == hashCode && (this.key == key || this.key.equals(key));
        }
    }
    
    static final class ListNode implements Node {
        
        private final Node[] table;

        public ListNode(Node[] table) {
            this.table = table;
        }
        
        @Override
        public Object get(final int hashCode, Object key) {
            for (int i = 0, n = table.length; i < n; i++) {
                Object value = table[i].get(hashCode, key);
                if (value != null) {
                    return value;
                }
            }
            return null;
        }

        @Override
        public ListNode put(int level, final int hashCode, Object key, Object value, Out<Object> out) {
            int index = find(hashCode, key, out);
            
            if (out.value == null) {
                return new ListNode(copyAndInsert(table, new LeafNode(hashCode, key, value), table.length));
            } else {
                return new ListNode(copyAndReplace(table, new LeafNode(hashCode, key, value), index));
            }
        }

        @Override
        public Node remove(final int hashCode, Object key, Out<Object> out) {
            int index = find(hashCode, key, out);
            
            if (out.value != null) {
                if (table.length != 2) {
                    return new ListNode(copyAndRemove(table, index));
                } else {                    
                    return table[index ^ 1];
                }
            }
            
            return this;
        }
        
        private int find(final int hashCode, Object key, Out<Object> out) {
            for (int i = 0, n = table.length; i < n; i++) {
                Object existingValue = table[i].get(hashCode, key);
                if (existingValue != null) {
                    out.value = existingValue;
                    return i;
                }
            }
            return 0;
        }
        
        public static ListNode create(Node nodeA, Node nodeB) {
            return new ListNode(new Node[] { nodeA, nodeB });
        }
    }    
    
    // ----------------- Utility Methods -------------------

    public static Node[] copyAndRemove(Node[] nodes, int index) {
        int len = nodes.length - 1;
        
        Node[] newNodes = new Node[len];
        arraycopy(nodes, 0, newNodes, 0, index);
        arraycopy(nodes, index + 1, newNodes, index, len - index);
        return newNodes;
    }

    public static Node[] copyAndReplace(Node[] nodes, Node node, int index) {
        int len = nodes.length;
        
        Node[] newNodes = new Node[len];
        System.arraycopy(nodes, 0, newNodes, 0, len);
        newNodes[index] = node;
        return newNodes;
    }

    public static Node[] copyAndInsert(Node[] nodes, Node node, int index) {
        int len = nodes.length;
    
        Node[] newNodes = (Node[]) new Node[len + 1];
        arraycopy(nodes, 0, newNodes, 0, index);
        arraycopy(nodes, index, newNodes, index + 1, len - index);
        newNodes[index] = node;
        return newNodes;
    }
}
